package com.bugpool.leilema.freamwork.mybatis;

import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import com.bugpool.leilema.freamwork.datascope.DataScopeConstants;
import com.bugpool.leilema.freamwork.datascope.DataScopeUtils;
import com.bugpool.leilema.tc.bean.entity.SysRole;
import com.bugpool.leilema.tc.bean.vo.SysUserVo;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.HexValue;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SetOperationList;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.io.StringReader;
import java.util.List;
import java.util.Properties;

import static com.bugpool.leilema.freamwork.datascope.DataScopeConstants.*;

/**
 * @Author wh
 * @Description mybatis 数据权限插件
 * @Date 2023/6/12 11:07
 * @Version 1.0
 **/
@Slf4j
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
})
public class MybatisDataPermissionIntercept1 implements Interceptor {
    //sql分析器
    CCJSqlParserManager parserManager = new CCJSqlParserManager();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        try {
            Object[] args = invocation.getArgs();
            MappedStatement ms = (MappedStatement) args[0];
            Object parameter = args[1];
            RowBounds rowBounds = (RowBounds) args[2];
            ResultHandler resultHandler = (ResultHandler) args[3];
            Executor executor = (Executor) invocation.getTarget();
            CacheKey cacheKey = null;
            BoundSql boundSql = null;
            //由于逻辑关系 只会进入一次
            if (args.length == 4) {
                boundSql = ms.getBoundSql(parameter);
                cacheKey = executor.createCacheKey(ms, parameter, rowBounds, boundSql);
            }

            if (args.length == 6) {
                cacheKey = (CacheKey) args[4];
                boundSql = (BoundSql) args[5];
            }
            //获取原始sql
            String sql = boundSql.getSql();
            log.info("原始sql {}", sql);
            //判断线程内 是否有权限信息
            String dataScope = DataScopeUtils.getDataScope();
            if (DataScopeConstants.DATA_SCOPE.equals(dataScope)) {
                //增强sql
                Select select = (Select) parserManager.parse(new StringReader(sql));
                SelectBody selectBody = select.getSelectBody();

                if (selectBody instanceof PlainSelect) {
                    this.setWhere((PlainSelect) selectBody);
                } else if (selectBody instanceof SetOperationList) {
                    SetOperationList setOperationList = (SetOperationList) selectBody;
                    List<SelectBody> selectBodyList = setOperationList.getSelects();
                    selectBodyList.forEach((s) -> {
                        this.setWhere((PlainSelect) s);
                    });
                }
                String dataPermissionSql = select.toString();
                log.info("增强SQL： {}", dataPermissionSql);
                BoundSql dataPermissionBoundSql = new BoundSql(ms.getConfiguration(), dataPermissionSql, boundSql.getParameterMappings(), parameter);
                return executor.query(ms, parameter, rowBounds, resultHandler, cacheKey, dataPermissionBoundSql);
            }

            return executor.query(ms, parameter, rowBounds, resultHandler, cacheKey, boundSql);

        } catch (Exception e) {

        }
        return null;
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        Interceptor.super.setProperties(properties);
    }


    protected void setWhere(PlainSelect plainSelect) {
        Expression sqlSegment = this.getSqlSegment(plainSelect.getWhere());
        if (sqlSegment!=null){
            plainSelect.setWhere(sqlSegment);
        }
    }

    @SneakyThrows
    public Expression getSqlSegment(Expression where) {
        SysUserVo loginUser = DataScopeUtils.getUser();
        if (loginUser == null){
            return where;
        }
        Long deptId = loginUser.getDeptId();
        Long userId = loginUser.getUserId();
        List<SysRole> sysRoleList = loginUser.getSysRoleList();
        StringBuilder sqlString = new StringBuilder();
        for (SysRole role : sysRoleList) {
            String dataScopeNum = role.getDataScope();
            Long roleId = role.getRoleId();
            if (DATA_SCOPE_ALL.equals(dataScopeNum)) {
                // 全部数据权限
                sqlString = new StringBuilder();
                break;
            } else if (DATA_SCOPE_CUSTOM.equals(dataScopeNum)) {
                sqlString.append(" OR dept_id IN ( SELECT dept_id FROM tb_sys_role_dept WHERE role_id = '")
                        .append(roleId)
                        .append("' ) ");
            } else if (DATA_SCOPE_DEPT.equals(dataScopeNum)) {
                sqlString.append(" OR dept_id = ").append(deptId).append(" ");
            } else if (DATA_SCOPE_DEPT_AND_CHILD.equals(dataScopeNum)) {
                sqlString.append(" OR dept_id IN ( WITH RECURSIVE dept_tree AS (\n" +
                        "  -- 查询本部门\n" +
                        "  SELECT dept_id FROM tb_sys_dept WHERE dept_id =")
                        .append(deptId)
                        .append("  UNION ALL\n" +
                                "  -- 查询本部门下的所有子部门\n" +
                                "  SELECT d.dept_id FROM dept_tree dt JOIN tb_sys_dept d ON dt.dept_id = d.parent_id\n" +
                                ")\n" +
                                "SELECT dept_id FROM dept_tree)");
            }else if (DATA_SCOPE_SELF.equals(dataScopeNum)) {
                //TODO 暂时有问题
                sqlString.append(" OR user_id = '").append(userId).append("' ");
            }
        }
        if (StringUtils.isNotBlank(sqlString.toString())) {
            if (where == null){
                where = new HexValue(" 1 = 1 ");
            }
            sqlString.insert(0," AND (");
            sqlString.append(")");
            sqlString.delete(7, 9);
            //判断是不是分页， 分页完成之后 清除权限标识
            return CCJSqlParserUtil.parseCondExpression(where + sqlString.toString());
        }else {
            return where;
        }
    }





}
